package org.zoomdev.zoom.aop.modules;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.zoomdev.zoom.aop.AopFactory;
import org.zoomdev.zoom.aop.factory.impl.AsmClassloader;
import org.zoomdev.zoom.aop.utils.AopUtils;
import org.zoomdev.zoom.aop.utils.AsmUtils;
import org.zoomdev.zoom.common.exceptions.ZoomException;
import org.zoomdev.zoom.common.res.ResScanner;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public final class AsmModuleContainer {

    private static ResScanner scanner;

    private static Map<Class, Map<String, MethodWithTypeConfig>> configMap = new HashMap<>();

    private static Map<Class, Class> enhancedMap = new ConcurrentHashMap<>();

    public static void register(Class moduleType, Map<String, MethodWithTypeConfig> config) {
        configMap.put(moduleType, config);
    }

    /**
     * @return
     */
    public static Class getEnhancedModule(Class moduleType, AopFactory factory) {
        return enhancedMap.computeIfAbsent(moduleType, k -> {
            Map<String, MethodWithTypeConfig> config = configMap.get(k);
            if (config == null) {
                return k;
            }
            int count = 0;
            for(Map.Entry<String,MethodWithTypeConfig> entry : config.entrySet()){
                MethodWithTypeConfig cfg =  entry.getValue();
                Class enhancedClass = factory.enhance(cfg.getRealType());
                if(enhancedClass!=cfg.getRealType()){
                    cfg.setEnhancedType(enhancedClass);
                    ++count;
                }
            }
            if(count==0){
                return k;
            }
            try {
                ResScanner.ClassRes res = scanner.getClass(k.getName());
                ClassReader cr = new ClassReader(res.getBytes());
                ClassWriter cw = new ClassWriter(cr, 0);
                String className = k.getName() + "$EnhanceModule";
                CopyMethodVisitor visitor = new CopyMethodVisitor(cw, AsmUtils.makeExternalName(className), config);
                cr.accept(visitor, 0);
                Class newType = AsmClassloader.defineClass(className, cw.toByteArray());
                return newType;
            } catch (IOException e) {
                throw new ZoomException(e);
            }

        });

    }

    public static ResScanner getScanner() {
        return scanner;
    }

    public static void setScanner(ResScanner scanner) {
        AsmModuleContainer.scanner = scanner;
    }
}
